package com.example.mcp;

import fi.iki.elonen.NanoHTTPD;
import io.modelcontextprotocol.json.McpJsonMapper;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpServerSession;
import io.modelcontextprotocol.spec.McpServerTransport;
import io.modelcontextprotocol.spec.McpServerTransportProvider;
import io.modelcontextprotocol.spec.ProtocolVersions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 基于 NanoHTTPD 的最小可编译版 MCP 服务器传输提供者骨架。
 *
 * 说明：当前仅提供占位实现，后续将补齐 SSE 与消息端点逻辑。
 */
public class NanoHttpdSseServerTransportProvider implements McpServerTransportProvider {

    private static final Logger logger = LoggerFactory.getLogger(NanoHttpdSseServerTransportProvider.class);

    public static final String DEFAULT_SSE_ENDPOINT = "/sse";
    public static final String DEFAULT_MESSAGE_ENDPOINT = "/message";

    private final int port;
    private final String sseEndpoint;
    private final String messageEndpoint;

    private McpServerSession.Factory sessionFactory;

    private final Map<String, SseSessionContext> sessions = new ConcurrentHashMap<>();
    private final AtomicBoolean isClosing = new AtomicBoolean(false);

    private NanoServer server;
    private final McpJsonMapper jsonMapper = createUtf8JsonMapper();
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

    private NanoHttpdSseServerTransportProvider(Builder builder) {
        this.port = builder.port;
        this.sseEndpoint = builder.sseEndpoint;
        this.messageEndpoint = builder.messageEndpoint;
    }

    public static Builder builder() {
        return new Builder();
    }

    /**
     * 创建支持UTF-8编码的JSON映射器
     */
    private static McpJsonMapper createUtf8JsonMapper() {
        try {
            // 创建自定义的Jackson ObjectMapper，确保UTF-8编码支持
            com.fasterxml.jackson.databind.ObjectMapper objectMapper = new com.fasterxml.jackson.databind.ObjectMapper();
            
            // 配置UTF-8编码
            objectMapper.configure(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES, true);
            objectMapper.configure(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_SINGLE_QUOTES, true);
            
            // 确保字符串处理使用UTF-8编码
            objectMapper.getFactory().configure(com.fasterxml.jackson.core.JsonGenerator.Feature.ESCAPE_NON_ASCII, false);
            
            // 创建自定义的McpJsonMapper
            return new io.modelcontextprotocol.json.jackson.JacksonMcpJsonMapper(objectMapper);
        } catch (Exception e) {
            logger.warn("Failed to create custom UTF-8 JSON mapper, using default: {}", e.getMessage());
            return McpJsonMapper.getDefault();
        }
    }

    /**
     * 手动修复编码问题
     */
    private String manuallyFixEncoding(String body) {
        try {
            // 尝试多种编码修复方式
            String[] encodings = {"UTF-8", "ISO-8859-1", "GBK", "GB2312"};
            for (String encoding : encodings) {
                try {
                    byte[] bytes = body.getBytes(java.nio.charset.StandardCharsets.UTF_8);
                    String fixed = new String(bytes, java.nio.charset.Charset.forName(encoding));
                    if (!fixed.contains("?")) {
                        logger.debug("使用编码 {} 修复成功: {}", encoding, fixed);
                        return fixed;
                    }
                } catch (Exception e) {
                    // 忽略编码错误，尝试下一个
                }
            }
            
            // 如果所有编码都失败，返回原始字符串
            logger.warn("所有编码修复方式都失败，返回原始字符串");
            return body;
        } catch (Exception e) {
            logger.error("手动修复编码失败: {}", e.getMessage());
            return body;
        }
    }

    @Override
    public void setSessionFactory(McpServerSession.Factory sessionFactory) {
        this.sessionFactory = sessionFactory;
        // 在设置后启动服务器，确保端点可用
        startServerIfNeeded();
    }

    @Override
    public Mono<Void> notifyClients(String method, Object params) {
        return Mono.fromRunnable(() -> {
            sessions.values().forEach(ctx -> {
                try {
                    ctx.session().sendNotification(method, params).block();
                } catch (Exception e) {
                    logger.warn("Failed to notify session {}: {}", ctx.sessionId(), e.toString());
                }
            });
        });
    }

    @Override
    public void close() {
        this.closeGracefully().subscribe();
    }

    @Override
    public Mono<Void> closeGracefully() {
        isClosing.set(true);
        return Mono.fromRunnable(() -> {
            try {
                if (server != null) {
                    server.stop();
                }
            } finally {
                try {
                    scheduler.shutdownNow();
                } catch (Exception ignored) {}
                sessions.clear();
            }
        });
    }

    @Override
    public List<String> protocolVersions() {
        return List.of(ProtocolVersions.MCP_2024_11_05);
    }

    private synchronized void startServerIfNeeded() {
        if (this.server != null) {
            return;
        }
        this.server = new NanoServer(this.port);
        try {
            this.server.start(NanoHTTPD.SOCKET_READ_TIMEOUT, false);
            logger.info("NanoHTTPD started on port {} (sse: {}, message: {})", port, sseEndpoint, messageEndpoint);
        } catch (IOException e) {
            throw new IllegalStateException("Failed to start NanoHTTPD server", e);
        }
    }

    public static class Builder {
        private int port = 8080;
        private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
        private String messageEndpoint = DEFAULT_MESSAGE_ENDPOINT;

        public Builder port(int port) {
            this.port = port;
            return this;
        }

        public Builder sseEndpoint(String sseEndpoint) {
            this.sseEndpoint = sseEndpoint;
            return this;
        }

        public Builder messageEndpoint(String messageEndpoint) {
            this.messageEndpoint = messageEndpoint;
            return this;
        }

        public NanoHttpdSseServerTransportProvider build() {
            return new NanoHttpdSseServerTransportProvider(this);
        }
    }

    /**
     * NanoHTTPD 服务器占位实现：后续将实现 /sse 与 /message 的交互与会话管理。
     */
    private class NanoServer extends NanoHTTPD {
        public NanoServer(int port) {
            super(port);
        }

        @Override
        public Response serve(IHTTPSession session) {
            if (isClosing.get()) {
                return NanoHTTPD.newFixedLengthResponse(Response.Status.SERVICE_UNAVAILABLE, "text/plain", "Server closing");
            }

            String uri = session.getUri();
            Method method = session.getMethod();

            if (Method.GET.equals(method) && uri.equals(sseEndpoint)) {
                return handleSse(session);
            }
            if (Method.POST.equals(method) && uri.equals(messageEndpoint)) {
                return handleMessage(session);
            }
            return NanoHTTPD.newFixedLengthResponse(Response.Status.NOT_FOUND, "text/plain", "Not Found");
        }

        private Response handleSse(IHTTPSession http) {
            String sessionId = UUID.randomUUID().toString();

            try {
                // 建立管道：将写端交给会话传输，读端作为响应体
                java.io.PipedOutputStream out = new java.io.PipedOutputStream();
                java.io.PipedInputStream in = new java.io.PipedInputStream(out, 16 * 1024);

                NanoSseMcpSessionTransport transport = new NanoSseMcpSessionTransport(sessionId, out);
                McpServerSession mcpSession = sessionFactory.create(transport);
                SseSessionContext ctx = new SseSessionContext(sessionId, mcpSession, transport, out);
                sessions.put(sessionId, ctx);

                // 初始 endpoint 事件
                String endpointUrl = this.getEndpointUrl(http, sessionId);
                ctx.sendEvent("endpoint", endpointUrl);

                Response r = NanoHTTPD.newChunkedResponse(Response.Status.OK, "text/event-stream", in);
                // SSE 标准头
                r.addHeader("Cache-Control", "no-cache");
                r.addHeader("Connection", "keep-alive");
                r.addHeader("Access-Control-Allow-Origin", "*");
                r.addHeader("Content-Type", "text/event-stream; charset=utf-8");
                return r;
            } catch (IOException e) {
                logger.error("Failed to init SSE: {}", e.toString());
                return NanoHTTPD.newFixedLengthResponse(Response.Status.INTERNAL_ERROR, "text/plain", "SSE init failed");
            }
        }

        private String getEndpointUrl(IHTTPSession http, String sessionId) {
            String hostHeader = http.getHeaders().getOrDefault("host", "localhost:" + port);
            String scheme = "http"; // NanoHTTPD 默认 http
            return scheme + "://" + hostHeader + messageEndpoint + "?sessionId=" + sessionId;
        }

        private Response handleMessage(IHTTPSession http) {
            Map<String, List<String>> params = http.getParameters();
            String sessionId = null;
            if (params != null && params.containsKey("sessionId") && !params.get("sessionId").isEmpty()) {
                sessionId = params.get("sessionId").get(0);
            }
            if (sessionId == null || sessionId.isBlank()) {
                return jsonError(Response.Status.BAD_REQUEST, new io.modelcontextprotocol.spec.McpError("Session ID missing in message endpoint"));
            }

            SseSessionContext ctx = sessions.get(sessionId);
            if (ctx == null) {
                return jsonError(Response.Status.NOT_FOUND, new io.modelcontextprotocol.spec.McpError("Session not found: " + sessionId));
            }

            try {
                // 直接读取请求体，确保UTF-8编码处理
                String body = "";
                try {
                    // 获取Content-Length
                    String contentLengthHeader = http.getHeaders().get("content-length");
                    int contentLength = contentLengthHeader != null ? Integer.parseInt(contentLengthHeader) : 0;
                    
                    if (contentLength > 0) {
                        // 直接读取原始字节，然后转换为UTF-8字符串
                        byte[] buffer = new byte[contentLength];
                        int bytesRead = 0;
                        int totalBytesRead = 0;
                        
                        while (totalBytesRead < contentLength) {
                            bytesRead = http.getInputStream().read(buffer, totalBytesRead, contentLength - totalBytesRead);
                            if (bytesRead == -1) {
                                break;
                            }
                            totalBytesRead += bytesRead;
                        }
                        
                        // 使用UTF-8编码将字节转换为字符串
                        body = new String(buffer, 0, totalBytesRead, java.nio.charset.StandardCharsets.UTF_8);
                    } else {
                        // 如果没有Content-Length，使用传统方式但确保UTF-8编码
                        Map<String, String> files = new ConcurrentHashMap<>();
                        http.parseBody(files);
                        body = files.get("postData");
                        if (body == null) body = "";
                    }

                } catch (Exception e) {
                    logger.warn("直接读取请求体失败，回退到parseBody方法: {}", e.getMessage());
                    // 回退到原始方法
                    Map<String, String> files = new ConcurrentHashMap<>();
                    http.parseBody(files);
                    body = files.get("postData");
                    if (body == null) body = "";
                }

                // 确保请求体使用UTF-8编码
                logger.info("Received request body: {}", body);
                
                // 如果body为空，返回错误
                if (body.trim().isEmpty()) {
                    logger.warn("Empty request body received");
                    return jsonError(Response.Status.BAD_REQUEST, new io.modelcontextprotocol.spec.McpError("Empty request body"));
                }
                
                // 添加详细的编码调试信息
                logger.info("Request body bytes: {}", java.util.Arrays.toString(body.getBytes(java.nio.charset.StandardCharsets.UTF_8)));
                logger.info("Request body charset: {}", java.nio.charset.Charset.defaultCharset());
                logger.info("Request body contains '?': {}", body.contains("?"));
                
                // 尝试修复可能的编码问题
                String fixedBody = body;
                if (body.contains("?")) {
                    logger.warn("检测到可能的编码问题，尝试修复...");
                    try {
                        // 尝试将ISO-8859-1编码的字符串转换为UTF-8
                        byte[] bytes = body.getBytes(java.nio.charset.StandardCharsets.ISO_8859_1);
                        fixedBody = new String(bytes, java.nio.charset.StandardCharsets.UTF_8);
                        logger.debug("修复后的body: {}", fixedBody);
                    } catch (Exception e) {
                        logger.warn("编码修复失败，使用原始body: {}", e.getMessage());
                        fixedBody = body;
                    }
                }
                
                // 手动解析JSON并修复编码问题
                McpSchema.JSONRPCMessage message;
                try {
                    // 先尝试正常解析
                    message = McpSchema.deserializeJsonRpcMessage(jsonMapper, fixedBody);
                } catch (Exception e) {
                    logger.warn("JSON解析失败，尝试手动修复编码: {}", e.getMessage());
                    // 如果解析失败，尝试手动修复编码
                    String manuallyFixedBody = manuallyFixEncoding(fixedBody);
                    message = McpSchema.deserializeJsonRpcMessage(jsonMapper, manuallyFixedBody);
                }
                // 阻塞处理，保持与 Servlet 实现一致
                ctx.session().handle(message).block();
                return NanoHTTPD.newFixedLengthResponse(Response.Status.OK, "application/json; charset=utf-8", "{}");
            } catch (Exception e) {
                logger.error("Error processing message: {}", e.toString());
                return jsonError(Response.Status.INTERNAL_ERROR, new io.modelcontextprotocol.spec.McpError("Error processing message: " + e.getMessage()));
            }
        }

        private Response jsonError(Response.Status status, io.modelcontextprotocol.spec.McpError error) {
            try {
                String json = jsonMapper.writeValueAsString(error);
                Response r = NanoHTTPD.newFixedLengthResponse(status, "application/json; charset=utf-8", json);
                r.addHeader("Access-Control-Allow-Origin", "*");
                return r;
            } catch (Exception e) {
                return NanoHTTPD.newFixedLengthResponse(status, "text/plain; charset=utf-8", error.getMessage());
            }
        }
    }

    private static String toSseEvent(String event, String data) {
        StringBuilder sb = new StringBuilder();
        if (event != null && !event.isEmpty()) {
            sb.append("event: ").append(event).append('\n');
        }
        // 将数据按行输出为 data: 行
        String[] lines = data.split("\\r?\\n");
        for (String line : lines) {
            sb.append("data: ").append(line).append('\n');
        }
        sb.append('\n');
        return sb.toString();
    }

    private static String toSseComment(String comment) {
        StringBuilder sb = new StringBuilder();
        String[] lines = comment.split("\\r?\\n");
        for (String line : lines) {
            sb.append(": ").append(line).append('\n');
        }
        sb.append('\n');
        return sb.toString();
    }

    private class SseSessionContext {
        private final String sessionId;
        private final McpServerSession session;
        private final NanoSseMcpSessionTransport transport;
        private final java.io.OutputStream out;
        private volatile ScheduledFuture<?> keepAliveFuture;

        SseSessionContext(String sessionId, McpServerSession session, NanoSseMcpSessionTransport transport, java.io.OutputStream out) {
            this.sessionId = sessionId;
            this.session = session;
            this.transport = transport;
            this.out = out;
        }

        String sessionId() { return sessionId; }
        McpServerSession session() { return session; }

        void writePrelude() throws IOException {
            // 约 2KB 注释填充，避免某些代理在初始阶段关闭空闲连接
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < 64; i++) {
                sb.append(": ");
                for (int j = 0; j < 30; j++) sb.append('x');
                sb.append('\n');
            }
            sb.append('\n');
            try {
                out.write(sb.toString().getBytes(java.nio.charset.StandardCharsets.UTF_8));
                out.flush();
            } catch (IOException e) {
                sessions.remove(this.sessionId);
                cancelKeepAlive();
                try { out.close(); } catch (IOException ignored) {}
                throw e;
            }
        }

        synchronized void sendEvent(String event, String data) throws IOException {
            String payload = toSseEvent(event, data);
            try {
                out.write(payload.getBytes(java.nio.charset.StandardCharsets.UTF_8));
                out.flush();
            } catch (IOException e) {
                // 客户端断开时可能触发管道错误，移除会话并吞掉异常
                sessions.remove(this.sessionId);
                cancelKeepAlive();
                try { out.close(); } catch (IOException ignored) {}
                throw e;
            }
        }

        void startKeepAlivePing() {
            cancelKeepAlive();
            this.keepAliveFuture = scheduler.scheduleAtFixedRate(() -> {
                try {
                    String payload = toSseComment("keep-alive " + System.currentTimeMillis());
                    synchronized (this) {
                        out.write(payload.getBytes(java.nio.charset.StandardCharsets.UTF_8));
                        out.flush();
                    }
                } catch (IOException ignored) {
                    // 已在 sendEvent 中清理
                }
            }, 15, 15, TimeUnit.SECONDS);
        }

        void cancelKeepAlive() {
            ScheduledFuture<?> f = this.keepAliveFuture;
            if (f != null) {
                f.cancel(true);
                this.keepAliveFuture = null;
            }
        }
    }

    private class NanoSseMcpSessionTransport implements McpServerTransport {
        private final String sessionId;
        private final java.io.OutputStream out;
        private final AtomicBoolean closed = new AtomicBoolean(false);

        NanoSseMcpSessionTransport(String sessionId, java.io.OutputStream out) {
            this.sessionId = sessionId;
            this.out = out;
        }

        @Override
        public Mono<Void> closeGracefully() {
            return Mono.fromRunnable(() -> {
                if (closed.compareAndSet(false, true)) {
                    try {
                        synchronized (this) {
                            String bye = toSseEvent("close", "closing");
                            out.write(bye.getBytes(java.nio.charset.StandardCharsets.UTF_8));
                            out.flush();
                        }
                    } catch (IOException ignored) {
                    } finally {
                        try { out.close(); } catch (IOException ignored) {}
                    }
                }
            });
        }

        @Override
        public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
            return Mono.fromRunnable(() -> {
                if (closed.get()) return;
                try {
                    String json = jsonMapper.writeValueAsString(message);
                    synchronized (this) {
                        String payload = toSseEvent("message", json);
                        out.write(payload.getBytes(java.nio.charset.StandardCharsets.UTF_8));
                        out.flush();
                    }
                } catch (Exception e) {
                    // 写失败通常意味着客户端断开，清理会话
                    sessions.remove(sessionId);
                    logger.warn("SSE write failed for session {}: {}", sessionId, e.toString());
                }
            });
        }

        @Override
        public <T> T unmarshalFrom(Object data, io.modelcontextprotocol.json.TypeRef<T> typeRef) {
            return jsonMapper.convertValue(data, typeRef);
        }
    }
}


